#pragma once

#include "webots/supervisor.h"
#include <ctrl_alg/pid/pid.hpp>
#include <ctrl_alg/pid/pid_pos.hpp>
#include <memory>
#include <rclcpp/rclcpp.hpp>
#include <webots_ros2_driver/PluginInterface.hpp>
#include <webots_ros2_driver/WebotsNode.hpp>

#include <general_interface/msg/chassis_servo.hpp>
#include <general_interface/msg/control_signal.hpp>

namespace webots_driver
{
class WebotsDriver : public webots_ros2_driver::PluginInterface
{
  private:
    using ControlSignal = general_interface::msg::ControlSignal;
    rclcpp::Subscription<ControlSignal>::SharedPtr sub;
    using ChassisServo = general_interface::msg::ChassisServo;
    rclcpp::Subscription<ChassisServo>::SharedPtr servoSub;
    WbDeviceTag rightWheel;
    WbDeviceTag leftWheel;
    WbDeviceTag rightSensor;
    WbDeviceTag leftSensor;
    WbDeviceTag imu;
    WbDeviceTag gyro;

    std::unique_ptr<CtrlAlg::PID<double>> gyro_pid;
    std::unique_ptr<CtrlAlg::PID<double>> angle_pid;
    std::unique_ptr<CtrlAlg::PID<double>> speed_pid;
    std::unique_ptr<CtrlAlg::PID<double>> rotation_pid;

    double lastLeftPosition, lastRightPosition;

    double forwardInput = 0;
    double rotationInput = 0;
    double forwardScale = 1;
    double rotationScale = 1;
    bool directControl = false;

    struct Servo
    {
        double angle = 0;
        double gyro = 0;
        double speed = 0;
        double gyroZ = 0;
        double angleGain = 1;
        double gyroGain = 0;
        double speedGain = 0;
        double gyroZGain = 0;
    };

    Servo target;

  public:
    void init(webots_ros2_driver::WebotsNode *node, std::unordered_map<std::string, std::string> &parameters) override;
    void step() override;

  private:
    std::unique_ptr<CtrlAlg::PIDPos<double>> createPID(const std::string &param);
    void OnChassisServo(const ChassisServo& servo);
    void DirectControl();
    void ChassisServoControl();
};
} // namespace webots_driver
